import numpy as np

def R_X(theta):
    """生成绕X轴旋转的量子门"""
    return np.array([
        [np.cos(theta / 2), -1j * np.sin(theta / 2)],
        [-1j * np.sin(theta / 2), np.cos(theta / 2)]
    ])

def R_Y(theta):
    """生成绕Y轴旋转的量子门"""
    return np.array([
        [np.cos(theta / 2), -np.sin(theta / 2)],
        [np.sin(theta / 2), np.cos(theta / 2)]
    ])

def R_Z(theta):
    """生成绕Z轴旋转的量子门"""
    return np.array([
        [np.exp(-1j * theta / 2), 0],
        [0, np.exp(1j * theta / 2)]
    ])

H0 = np.array([
    [1/np.sqrt(2), 1/np.sqrt(2)],
    [1/np.sqrt(2), -1/np.sqrt(2)]
])
Y2M=R_Y(-np.pi/2)
Y2P=R_Y(np.pi/2)

# 定义矩阵 M1
SWAP = np.array([
    [1, 0, 0, 0],
    [0, 0, 1, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1]
])
ISWAP = np.array([
    [1, 0, 0, 0],
    [0, 0, 1j, 0],
    [0, 1j, 0, 0],
    [0, 0, 0, 1]
])

# 定义矩阵 M2
CNOT = np.array([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 0, 1],
    [0, 0, 1, 0]
])

# 定义矩阵 M3
CZ = np.array([
    [1, 0, 0, 0],
    [0, 1, 0, 0],
    [0, 0, 1, 0],
    [0, 0, 0, -1]
])
I=np.array([
    [1, 0],
    [0, 1]
])

def CRZ(phi):
    """生成控制Z门"""
    return np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, np.exp(1j * phi)]
    ])

def CRX(phi):
    """生成控制X门"""
    return np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, np.cos(phi / 2), -1j * np.sin(phi / 2)],
        [0, 0, -1j * np.sin(phi / 2), np.cos(phi / 2)]
    ])

def CRY(phi):
    """生成控制Y门"""
    return np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, np.cos(phi / 2), -np.sin(phi / 2)],
        [0, 0, np.sin(phi / 2), np.cos(phi / 2)]
    ])
print('----------check Haradamard gate----------')
print('H_true:',H0)
print(np.dot(R_Z(np.pi),R_Y(np.pi/2)))

print('----------check CNOT gate----------')
print('CNOT_true:',CNOT)
print('CNOT_false:',np.dot(np.dot(np.kron(I,Y2P),CZ),np.kron(I,Y2M)))
print(np.allclose(np.dot(np.dot(np.kron(I,Y2P),CZ),np.kron(I,Y2M)) ,CNOT))

print('----------check SWAP gate----------')
# 第一组操作
step1 = np.kron(np.eye(2), Y2P) @ CZ @ np.kron(np.eye(2), Y2M)

# 第二组操作
step2 = np.kron(Y2P, np.eye(2)) @ CZ @ np.kron(Y2M, np.eye(2))

# 第三组操作
step3 = np.kron(np.eye(2), Y2P) @ CZ @ np.kron(np.eye(2), Y2M)

# 总的 SWAP 分解
SWAP_decomposed = step1 @ step2 @ step3

SWAP_original = SWAP
print("分解后的 SWAP 矩阵:\n", SWAP_decomposed)
print("原 SWAP 矩阵:\n", SWAP_original)
print("分解后的 SWAP 矩阵与原 SWAP 矩阵是否相等:", np.allclose(SWAP_decomposed, SWAP_original))